'''
完成适配，完成注释
参考链接：
1. [deepseek多模态模型（janus）微调实战。_哔哩哔哩_bilibili](https://www.bilibili.com/video/BV1r5AUeBEtk/?spm_id_from=333.1387.favlist.content.click&vd_source=4d2a71ad349005831356ddd323e9568b)
2. https://github.com/modelscope/ms-swift
3. https://modelscope.cn/datasets/AI-ModelScope/LaTeX_OCR
'''

import torch
from transformers import AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from modelscope import snapshot_download
from modelscope import MsDataset

# 定义模型路径，
model_path = r"M:\Projects\python\unsloth\janus_7b\Janus-Pro-7B"
print("模型路径：", model_path)
img_path = r"M:\Projects\python\unsloth\janus_7b\image\2.png"
print("图片路径：", img_path)
question = "图片是Latex的图片，请输出图片中的内容，格式为Markdown格式。"
conversation = [
    {
        "role": "<|User|>",
        "content": f"<image_placeholder>\n{question}\n",
        "images": [img_path]
    },
    {
        "role": "<|Assistant|>",
        "content": ""
    }
]


def load_latex_ocr_dataset():
    train_dataset = MsDataset.load("AI-ModelScope/LaTeX_OCR", subset_name="small", split="train")
    print(train_dataset[0])
    return train_dataset


def test_model(vl_chat_processor, vl_gpt, tokenizer):
    # 传入一个结构化的对话数据，是一个列表，包含多个对话轮次
    # 每个轮次是一个字典，包含 role、content、images 等键
    # images 键包含该轮次相关的图像路径列表
    # load_pil_images 遍历 [conversation]janus_7b.py ) 中的每个对话轮次，寻找包含 "images" 键的条目，将其中的图片进行加载到内存中
    # 返回一系列图片的列表：返回: List[PIL.Image.Image] 顺序：按照在 [conversation]janus_7b.py ) 中出现的顺序
    # 作为后续处理的输入
    # 如果图片过大，可以在这里加载完成后进行缩放处理，因为加载的PIL图片是原始的，可能会占用较多内存
    pil_images = load_pil_images(conversation) 
    '''
    图像预处理：调整尺寸、归一化、转换为张量
    文本预处理：分词、编码、构建注意力掩码
    格式标准化：将对话转换为模型训练时的格式

    输出包含：
    input_ids：文本的 token ID 序列
    attention_mask：注意力掩码
    pixel_values：预处理后的图像张量
    sft_format：格式化的对话文本
    '''
    prepare_inputs = vl_chat_processor(
        conversations=conversation,
        images=pil_images,
        force_batchify= True, # 确保输出是批处理格式
    ).to(vl_gpt.device)
    # 如果显存不足对模型进行转换为float16、bfloat16之类的，这里输入的数据
    # 也要进行转换，否则会报错输入的数据和实际的模型数据类型不匹配
    # 通过测试kimiv2 和 claude 3.7 sonnet 都只考虑到了模型的数据类型，没有考虑输入数据的类型
    # claude 4 sonnet 考虑到了输入数据的类型，但是转换的方式是错误的
    # 网页版 o3 模型 + web search 完美的解决了
    prepare_inputs.pixel_values = prepare_inputs.pixel_values.to(dtype=torch.float16)

    #  准备输入嵌入（融合图像和文本），将图像和文本进行融合为一个嵌入
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    # 使用融合后的嵌入进行文本生成
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id = tokenizer.eos_token_id, # 标识填充的token
        bos_token_id = tokenizer.bos_token_id, # 标识开始的token
        eos_token_id = tokenizer.eos_token_id, # 标识结束的token
        max_new_tokens=512, # 最大生成长度
        do_sample=False,  # 使用贪婪解码
        use_cache=True, # 启用KV缓存加速
    )

    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    print(f"{prepare_inputs['sft_format'][0]}: {answer}")


def main():
    model_file = snapshot_download("deepseek-ai/Janus-Pro-7B", local_dir=model_path)

    # Janus 模型的多模态数据预处理器，它的主要职责是将原始的图像和文本数据转换为模型能够理解的格式
    vl_chat_processor = VLChatProcessor.from_pretrained(model_file)
    tokenizer = vl_chat_processor.tokenizer
    
    # ugging Face Transformers 库的自动模型加载器，在这里它加载的是 Janus 的多模态因果语言模型
    '''
    视觉编码器 (vision_model)：基于 SigLIP，处理图像
    语言模型 (language_model)：基于 DeepSeek，处理文本
    对齐器 (aligner)：连接视觉和语言表示
    '''
    vl_gpt = AutoModelForCausalLM.from_pretrained(
        model_file, 
        trust_remote_code=True, # 允许自行自定义代码
        torch_dtype=torch.float16)
    '''
    trust_remote_code意思就是模型的文件内部会有自定义的代码，类似java的动态加载吧
    这样应该就可以兼容不同的模型，而不用等着transformer官方去更新，只需要实现标准的接口
    而如果模型的结构和Transformer相同的话，那么就先不要自定义的模块了，完全用标准的组件
    就不需要。
    一般是针对那些不明模型最好做限制，如果是大厂的开源模型，应该是可以放心使用的
    '''
    
    if torch.cuda.is_available():
        print("CUDA可用，使用GPU进行推理")
        vl_gpt = vl_gpt.cuda().eval()

    # 对齐所有模型的数据类型
    vl_gpt.vision_model.to(dtype=torch.float16)

    test_model(vl_chat_processor, vl_gpt, tokenizer)

if __name__ == "__main__":
    main()
